#ifndef MYAMG_SOLVERKERNELS_CUH
#define MYAMG_SOLVERKERNELS_CUH
#include <Data/CSRMatrixOperator.cuh>

namespace myamg {
    template<typename Integer, typename Real>
    void calcResidualHost(Real *r, const Integer *Ap, const Integer *Ai, const Real *Av, const Real *b, const Real *x,
                          const size_t n) {
#pragma omp parallel for
        for (auto i = 0; i < n; ++i) {
            auto temp = b[i];
            for (auto j = Ap[i]; j < Ap[i + 1]; ++j) {
                temp -= Av[j] * x[Ai[j]];
            }
            r[i] = temp;
        }
    }

    template<typename Integer, typename Real>
    __global__ void calcResidualDevice(Real *r, const Integer *Ap, const Integer *Ai, const Real *Av, const Real *b, const Real *x,
                          const size_t n) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < n) {
            auto temp = b[idx];
            for (auto j = Ap[idx]; j < Ap[idx + 1]; ++j) {
                temp -= Av[j] * x[Ai[j]];
            }
            r[idx] = temp;
        }
    }

    template<typename Real>
    void updatePHost(Real* p, const Real* r, const Real* v, const Real omega, const Real beta, const size_t n) {
#pragma omp parallel for
        for (auto i = 0; i < n; ++i) {
            p[i] = r[i] - beta * omega * v[i] + beta * p[i];
        }
    }

    template<typename Real>
    __global__ void updatePDevice(Real* p, const Real* r, const Real* v, const Real omega, const Real beta, const size_t n) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < n) {
            p[idx] = r[idx] - beta * omega * v[idx] + beta * p[idx];
        }
    }

    template<typename Real>
    void BiCGStabUpdateRHost(Real* r, const Real* v, const Real alpha, const size_t n) {
#pragma omp parallel for
        for (auto i = 0; i < n; ++i) {
            r[i] = r[i] - alpha * v[i];
        }
    }

    template<typename Real>
    __global__ void BiCGStabUpdateRDevice(Real* r, const Real* v, const Real alpha, const size_t n) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < n) {
            r[idx] = r[idx] - alpha * v[idx];
        }
    }

    template<typename Real>
    void BiCGStabUpdateXHost(Real* x, const Real* q, const Real* s, const Real alpha, const Real omega, const size_t n) {
#pragma omp parallel for
        for (auto i = 0; i < n; ++i) {
            x[i] = x[i] + alpha * q[i] + omega * s[i];
        }
    }

    template<typename Real>
    __global__ void BiCGStabUpdateXDevice(Real* x, const Real* q, const Real* s, const Real alpha, const Real omega, const size_t n) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < n) {
            x[idx] = x[idx] + alpha * q[idx] + omega * s[idx];
        }
    }
}

#endif //MYAMG_SOLVERKERNELS_CUH
